ENH: testing.lazy_xp_function: `torch.compile` support by lucascolley · Pull Request #668 · data-apis/array-api-extra

lucascolley · 2026-04-03T11:59:20Z

@rgommers this was easier than I expected. Is there someone familiar with Dynamo that we could ping? A few things that came up:

The following deprecation warning with a confusing message given that we are already using torch.compile:

E           torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
E           DeprecationWarning: `torch.jit.script_method` is not supported in Python 3.14+ and may break. Please switch to `torch.compile` or `torch.export`.

Some failing tests here with:

E                       torch._dynamo.exc.FailOnRecompileLimitHit: recompile_limit reached with fullgraph=True. Excessive recompilations can degrade performance due to the compilation overhead of each recompilation. To monitor recompilations, enable TORCH_LOGS=recompiles. If recompilations are expected, consider increasing torch._dynamo.config.cache_size_limit to an appropriate value.

(Draft PR as this will need docs updates before merge.)

lucascolley · 2026-04-03T12:01:19Z

+        if jit_library is JitLibrary.jax and isinstance(obj, Iterator):
            self._obj = list(obj)
            self._is_iter = True


claude reckoned that, unlike JAX, we needn't treat iterables with this special case for torch.compile

lucascolley · 2026-04-03T12:02:15Z

+                wrapped = autojit(func, JitLibrary.jax)
+                # If we're dealing with a staticmethod or classmethod, make
+                # sure things stay that way.
+                if isinstance(attr, staticmethod):
+                    wrapped = staticmethod(wrapped)
+                elif isinstance(attr, classmethod):
+                    wrapped = classmethod(wrapped)
+                temp_setattr(target, name, wrapped)
+
+    elif is_torch_namespace(xp):
+        for target, name, attr, func, tags in iter_tagged():
+            if tags["torch_compile"]:
+                wrapped = autojit(func, JitLibrary.torch)


(minor: could reduce some LoC perhaps)

lucascolley · 2026-04-03T12:02:41Z

@ev-br FYI

lucascolley

I think the warnings about to flood CI should be fixed by data-apis/array-api-compat#411 (comment).

lucascolley · 2026-04-03T13:58:12Z

ready for review

j-bowhay · 2026-08-16T07:03:03Z

@lucascolley what's the current status here? Is it just review bandwidth preventing this from moving forward?

lucascolley · 2026-08-16T10:23:22Z

@lucascolley what's the current status here? Is it just review bandwidth preventing this from moving forward?

Yes, pretty much. I had also noted in the linked issue:

Currently unclear how to deal with accepting a variable amount of graph breaks / recompilations.

ev-br · 2026-08-16T12:47:17Z

Currently unclear how to deal with accepting a variable amount of graph breaks / recompilations.

While I haven't written a single LOC, a reasonable design looks similar what is done for dask: an extra num_graphbreaks : int parameter (a made-up name), where a zero valueal maps onto fullgraph=True.

Also we definitely want to test things twice: in the eager mode and under torch dynamo. Not sure if it's already done for Jax (and if yes, where) or if it needs implementing.

j-bowhay · 2026-08-16T12:48:31Z

@lucascolley what's the current status here? Is it just review bandwidth preventing this from moving forward?

Yes, pretty much. I had also noted in the linked issue:

Currently unclear how to deal with accepting a variable amount of graph breaks / recompilations.

Do we need that level of granularity? Could we not have 3 options:

No support for torch.compile
Support for torch.compile but with full_graph=False
Support for torch.compile but with full_graph=True

lucascolley · 2026-08-16T12:53:17Z

@lucascolley what's the current status here? Is it just review bandwidth preventing this from moving forward?

Yes, pretty much. I had also noted in the linked issue:

Currently unclear how to deal with accepting a variable amount of graph breaks / recompilations.

Do we need that level of granularity? Could we not have 3 options:
* No support for `torch.compile`

* Support for `torch.compile` but with `full_graph=False`

* Support for `torch.compile` but with `full_graph=True`

Would be happy to start with that at least if we can get that working.

j-bowhay · 2026-08-16T13:10:19Z

Do you have a preference for what that api would look like? Would it be acceptable to have both a torch_compile and torch_full_graph argument?

lucascolley · 2026-08-16T13:17:00Z

Do you have a preference for what that api would look like? Would it be acceptable to have both a torch_compile and torch_full_graph argument?

not sure yet what I think is nicest... see also #917

j-bowhay · 2026-08-16T13:23:59Z

Do you have a preference for what that api would look like? Would it be acceptable to have both a torch_compile and torch_full_graph argument?

not sure yet what I think is nicest... see also #917

I guess the one difference being we also want to be able to turn compile on/off and control if graph breaks are allowed

lucascolley · 2026-08-17T01:34:01Z

CI is happy now, with fullgraph=False. TODOs:

ENH: fix is_jax_array under torch.compile array-api-compat#464
investigate the warnings that are ignored here
figure out the API for fullgraph=True
add all three modes (eager, no fullgraph, fullgraph) to array-api-extra's CI matrix
investigate recompilations / graph breaks in array-api-extra and perhaps remove some of them
add docs
try to plumb this into SciPy

lucascolley added lazy arrays xpx.testing labels Apr 3, 2026

lucascolley mentioned this pull request Apr 3, 2026

ENH: testing.lazy_xp_function: torch.compile support? #664

Open

lucascolley commented Apr 3, 2026

View reviewed changes

lucascolley requested a review from rgommers April 3, 2026 12:03

lucascolley commented Apr 3, 2026

View reviewed changes

Comment thread src/array_api_extra/_lib/_helpers.py

lucascolley commented Apr 3, 2026

View reviewed changes

Comment thread pyproject.toml Outdated

lucascolley force-pushed the torch-autojit branch 2 times, most recently from 393dd61 to cc4a106 Compare April 3, 2026 13:46

This was referenced Apr 3, 2026

ENH: array_namespace: support torch.compile data-apis/array-api-compat#413

Merged

ENH: torch.compile support scipy/scipy#24972

Open

lucascolley force-pushed the torch-autojit branch from cc4a106 to f16d510 Compare May 24, 2026 23:42

lucascolley force-pushed the torch-autojit branch from f16d510 to 27d1036 Compare August 16, 2026 13:13

ENH: testing.lazy_xp_function: torch.compile support

3578053

lucascolley force-pushed the torch-autojit branch from 27d1036 to 3578053 Compare August 16, 2026 14:00

TST: searchsorted: widen NaN sorting skip for torch

1e65e14

lucascolley force-pushed the torch-autojit branch 2 times, most recently from 03c9eb8 to b6cc704 Compare August 16, 2026 14:44

fix test paths

de01515

TST: autojit: restrict object inequality check to JAX

20003c6

lucascolley force-pushed the torch-autojit branch from b6cc704 to 20003c6 Compare August 16, 2026 14:48

add compiler dep

1297ba0

Uh oh!

Conversation

lucascolley commented Apr 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

lucascolley Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

lucascolley Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

lucascolley commented Apr 3, 2026

Uh oh!

Uh oh!

lucascolley left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

lucascolley commented Apr 3, 2026

Uh oh!

j-bowhay commented Aug 16, 2026

Uh oh!

lucascolley commented Aug 16, 2026

Uh oh!

ev-br commented Aug 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

j-bowhay commented Aug 16, 2026

Uh oh!

lucascolley commented Aug 16, 2026

Uh oh!

j-bowhay commented Aug 16, 2026

Uh oh!

lucascolley commented Aug 16, 2026

Uh oh!

j-bowhay commented Aug 16, 2026

Uh oh!

lucascolley commented Aug 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

lucascolley commented Apr 3, 2026 •

edited

Loading

ev-br commented Aug 16, 2026 •

edited

Loading